from pyais import decode as ais_decode, NMEAMessage
from pyais.exceptions import MissingMultipartMessageException


class AisDecoder:
    def __init__(self):
        self.fragment_cache = {}
        self.fragment_timeout = 3

    def decode_with_cache(self, name: str):
        try:
            ais_data = ais_decode(name)
            return ais_data.asdict()
        except MissingMultipartMessageException as e:
            fragment_ais_data = NMEAMessage.decode(NMEAMessage.from_string(name))
            key = fragment_ais_data.mmsi
            total_fragments = name.split(',')[1]
            if key not in self.fragment_cache:
                self.fragment_cache[key] = []
            self.fragment_cache[key].append(name)
            if len(self.fragment_cache[key]) == int(total_fragments):
                ais_data = ais_decode(self.fragment_cache[key])
                self.fragment_cache.pop(key)
                return ais_data.asdict()
